
import utils
import parameters
import os
import random
from paddle.io import Dataset


#将一个文件夹的图片文件切割，并存到指定文件夹中
def segImages(image_path,to_path,h=96,w=96):

    if not os.path.exists(to_path):
        os.makedirs(to_path)
    paths = os.listdir(image_path)
    idx = 0
    block = 0
    end = len(paths)
    
    for p in paths:
        idx += 1
        img = utils.read_image_rgb(os.path.join(image_path,p))
        segs = utils.seg_img(img,h,w)
        block = 0
        for seg in segs:
            block+=1
            utils.save_rgb(seg,os.path.join(to_path,"%d_%d.jpg"%(idx,block)))
        utils.print_progress_bar(idx,end)


#数据集加载类
class AnimeDataset(Dataset):
    
    def __init__(self,dataset,scale=2):
        
        super(AnimeDataset, self).__init__()
        self.dataset = dataset
        self.scale = scale
        self.img_path = os.listdir(dataset)
        self.num_samples = len(self.img_path)

    def __getitem__(self, index):
        
        img = utils.read_image_rgb(os.path.join(self.dataset,self.img_path[index]))
        h,w,_ = img.shape
        label = random.choice(utils.flip(img))  #随机选择增强后的一张图返回
        data = utils.resize_bicubic(label,h//self.scale,w//self.scale)
        label = utils.channel_last_to_first(label)
        data = utils.channel_last_to_first(data)
        #数据归一化到0-1
        return data/255., label/255

    def __len__(self):
        
        return self.num_samples
    
    
if __name__ == "__main__":
    if parameters.seg_train:
        print("seg training sets:")
        print("\nx2:")
        segImages(parameters.train_picture,parameters.train_dataset+"/x2",h=96,w=96)
        print("\nx3:")
        segImages(parameters.train_picture,parameters.train_dataset+"/x3",h=144,w=144)
        print("\nx4:")
        segImages(parameters.train_picture,parameters.train_dataset+"/x4",h=192,w=192)
    if parameters.seg_val:
        print("seg evaluate sets:")
        print("\nx2:")
        segImages(parameters.val_picture,parameters.val_dataset+"/x2",h=96,w=96)
        print("\nx3:")
        segImages(parameters.val_picture,parameters.val_dataset+"/x3",h=144,w=144)
        print("\nx4:")
        segImages(parameters.val_picture,parameters.val_dataset+"/x4",h=192,w=192)
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        